# -*- coding: utf-8 -*-
# File name: problem1.py

from projectq.backends import Simulator
from projectq.ops import QubitOperator, TimeEvolution, X, Z, H, Measure, All
from projectq.cengines import ( MainEngine,
                                AutoReplacer,
                                LocalOptimizer,
                                TagRemover, 
                                InstructionFilter,
                                DecompositionRuleSet)
import  projectq.setups.decompositions

import time
from scipy.sparse.linalg import eigs

# from mpi4py import MPI
import numpy as np


def get_hamiltonian(N):

    SumH0 = -0.5 * QubitOperator('X0')
    for i in range(1,N):
        SumH0 += -0.5 * QubitOperator("X" + str(i))

    SumHi = -0.5 * QubitOperator("X0")
    for i in range(1,N):
        SumHi += -0.5 * QubitOperator("X" + str(i))

    SumHi += -1 * QubitOperator("Z0")
    for i in range(1,N):
        SumHi += -1 * QubitOperator("Z" + str(i))


    SumHi += -2.0 * QubitOperator("X0 X1")
    for i in range(1,N-1):
        SumHi += -2.0 * QubitOperator(r"X%d X%d"%(i,i+1))

    SumHi += -2.0 * QubitOperator("Z0 Z1")
    for i in range(1,N-1):
        SumHi += -2.0 * QubitOperator(r"Z%d Z%d"%(i,i+1))

    return SumH0, SumHi


def get_kron_hamiltonian(N, alpha, beta, Jexc):
    X = np.array([[0, 1], [1, 0]])
    Z = np.array([[1, 0], [0, -1]])
    I = np.array([[1, 0], [0, 1]])

    Hamiltonian = -alpha*np.kron(X, np.eye(2**(N-1)))
    for j in range(1,N-1):
        a = np.eye(2**j)
        b = np.eye(2**(N-j-1))
        Hamiltonian += -alpha*np.kron(np.kron(a,X), b)
    Hamiltonian += -alpha*np.kron(np.eye(2**(N-1)), X)

    Hamiltonian += -beta*np.kron(Z, np.eye(2**(N-1)))
    for j in range(1,N-1):
        a = np.eye(2**j)
        b = np.eye(2**(N-j-1))
        Hamiltonian += -beta*np.kron(np.kron(a,Z), b)
    Hamiltonian += -beta*np.kron(np.eye(2**(N-1)), Z)

    Hamiltonian += -Jexc*np.kron(np.kron(X,X), np.eye(2**(N-2)))
    for j in range(1,N-3):
        a = np.eye(2**j)
        b = np.eye(2**(N-j-2))
        Hamiltonian += -Jexc*np.kron(np.kron(a,np.kron(X,X)), b)
    Hamiltonian += -Jexc*np.kron(np.eye(2**(N-2)), np.kron(X,X))

    Hamiltonian += -Jexc*np.kron(np.kron(Z,Z), np.eye(2**(N-2)))
    for j in range(1,N-3):
        a = np.eye(2**j)
        b = np.eye(2**(N-j-2))
        Hamiltonian += -Jexc*np.kron(np.kron(a,np.kron(Z,Z)), b)
    # Hamiltonian += -Jexc*np.kron(np.eye(2**(N-2)), np.kron(Z,Z))

    return Hamiltonian
                
    
def adiabatic_simulation():
    """
    Returns:
        real_energy(float):
            The final ideally continously evolved energy.
        simulated_energy(float):
            The final energy simulated by your mode1.
    """
    N = 13
    quregs = eng.allocate_qureg(N)
    All(H) | quregs
    eng.flush()

    SumH0, SumHi = get_hamiltonian(N)
    DeltaT = 3.0/20

    E0 = eng.backend.get_expectation_value(SumH0, quregs)

    for i in range(0,21):
        s = i/20.0
        SumH = (1 - s) * SumH0 + s * SumHi
        TimeEvolution(DeltaT, SumH) | quregs
        eng.flush()

    simulated_energy =  eng.backend.get_expectation_value(SumH, quregs)

    # print("Ground state energy: %f"%E0)
    # print("Simulated energy: %f"%simulated_energy)

    ####### Get matrix of Hi #######
    value0 = get_kron_hamiltonian(N, 0.5, 1.0, 2.0)
    ######Only calculate smallest eigen value#######
    a, b = eigs(value0, k=1, which='SR')
    real_energy = a[0].real
    # print("Approximate real energy: %f"%real_energy)
    All(Measure) | quregs

    return simulated_energy-E0, real_energy-E0
    

if __name__ == "__main__":
    # backend = SimulatorMPI(gate_fusion=True)
    
    # eng = HiQMainEngine(backend, engines)

    cache_depth = 10
    rule_set = DecompositionRuleSet(modules=[projectq.setups.decompositions])
    engines = [TagRemover()
                , LocalOptimizer(cache_depth)
                , AutoReplacer(rule_set)
                , TagRemover()
                , LocalOptimizer(cache_depth)
                # , GreedyScheduler()
              ]

    # Get init state
    eng = MainEngine(Simulator(), engines)

    # make the complier and run the circuit on the simulator backend

    simulated_energy, real_energy = adiabatic_simulation()

    simulated_error = (simulated_energy - real_energy)/real_energy
    simulated_error = np.abs(simulated_error)

    print(simulated_error)

    